#lang typed/racket
(require racket/unsafe/ops)

(: date-bytes ( -> Bytes))
(define (date-bytes)
  (define tstamp (quotient (current-seconds) 200))
  (list->bytes
   (list (remainder tstamp 256)
         (remainder (quotient tstamp 256) 256)
         (remainder (quotient (quotient tstamp 256) 256) 256)
         (remainder (quotient (quotient (quotient tstamp 256) 256) 256) 256))))

(: swap! ((Boxof Integer) (Boxof Integer) -> Void))
(define (swap! ba bb)
  (define temp (unbox ba))
  (set-box! ba (unbox bb))
  (set-box! bb temp))

(: make-rc4 (Bytes -> (Integer -> Integer)))
(define (make-rc4 keyd)
  
  ;New version per paper; disabled for compatibility
  ;(define key (make-bytes 128))
  ;(define tstamp (date-bytes))
  ;(for ([i (bytes-length keyd)])
  ;  (bytes-set! key i (bytes-ref keyd i)))
  ;(for ([i 4])
  ;  (bytes-set! key (+ i (bytes-length keyd))
  ;              (bytes-ref tstamp i)))
  
  (define key keyd)
  
  (define S
    (make-bytes 256 0))
  (for ([i 256])
    (bytes-set! S i i))
  (define j 0)
  (for ([i 256])
    (set! j
          (unsafe-fxmodulo
           (unsafe-fx+ j 
                       (unsafe-fx+ (bytes-ref S i)
                                   (bytes-ref key (unsafe-fxmodulo i (unsafe-bytes-length key)))))
           256))
    (define tmp (unsafe-bytes-ref S i))
    (unsafe-bytes-set! S i (bytes-ref S j))
    (unsafe-bytes-set! S j tmp))
  
  (set! j 0)
  (define i 0)
  
  (: Sr (Integer -> Integer))
  (define (Sr x) (unsafe-bytes-ref S x))
  
  (: toret (Integer -> Integer))
  (define (toret c)
    (set! i (unsafe-fxmodulo (add1 i) 256))
    (set! j (unsafe-fxmodulo (+ j (Sr i)) 256))
    (define temp (Sr i))
    (unsafe-bytes-set! S i (Sr j))
    (unsafe-bytes-set! S j temp)
    (bitwise-xor (Sr (unsafe-fxmodulo (+ (Sr i) (Sr j)) 256)) c))
  toret)

(: correct-key? (Bytes Bytes Bytes -> Boolean))
(define (correct-key? ctxt ptxt key)
  (define crypter (make-rc4 key))
  (: aux (Integer -> Boolean))
  (define (aux i)
    (cond
      [(= i (sub1 (bytes-length ctxt)))
       (printf "YAY CORRECT KEY ~a ~a ~a\n" ctxt ptxt key)
       #t]
      [else (and (= (crypter (bytes-ref ctxt i))
                    (bytes-ref ptxt i))
                 (aux (add1 i)))]))
  (aux 0))

;(correct-key? ((make-rc4 #"\0\0") #"i love you") #"i love you" #"\0\0")

(define ctr 0)

(: rc4-bruteforce (Bytes Bytes Bytes -> Bytes))
(define (rc4-bruteforce ctxt ptxt key)
  (when (= 0 (modulo (begin (set! ctr (add1 ctr)) ctr) 1000))
    (printf "~a: ~s\n" ctr key))
  (cond
    [(correct-key? ctxt ptxt key) key]
    [else (for ([i (bytes-length key)])
            (bytes-set! key i (random 256)))
          (rc4-bruteforce ctxt ptxt key)]))

(: rc4-findkey (Bytes Bytes -> Bytes))
(define (rc4-findkey ptxt ctxt)
  (rc4-bruteforce ctxt ptxt (make-bytes 2)))

(: arraychugger ((Integer -> Integer) -> (Bytes Integer -> Bytes)))
(define (arraychugger rc4chugger)
  (define (toret bt len)
    (define tr (make-bytes (bytes-length bt)))
    (for ([i len])
      (bytes-set! tr i (rc4chugger (bytes-ref bt i))))
    tr)
  toret)

;(define chug (arraychugger (make-rc4 #"cs")))

;(rc4-findkey #"i love you" (chug #"i love you"))

(provide (all-defined-out))